# -*- coding: utf-8 -*-
"""
Created on Fri Jul 12 16:21:00 2024
@author: xinyuan.wei
"""
import os
import pandas as pd
import rasterio
from pyproj import Transformer

# Get the current working directory
current_directory = os.getcwd()

# Get the parent directory
parent_directory = os.path.dirname(current_directory)

# Read the FIA data
FIA_file = os.path.join(parent_directory, 'NE_Plot_Composition.csv')
FIAdata = pd.read_csv(FIA_file)

# Filter the data based on the given conditions
FIAdata = FIAdata[(FIAdata['CONDID'] == 1) & 
                  (FIAdata['CONDPROP_UNADJ'] >= 0.9) &
                  (FIAdata['eAG'] >= 1) &
                  (FIAdata['STDAGE'] > 0) &
                  (FIAdata['STDAGE'] <= 200)].copy()

# Read the TIF climate data
pparent_directory = os.path.dirname(parent_directory)
climate_dir = os.path.join(pparent_directory, '3.Climate Data', 'ClimateNA')

# Initialize a transformer to ensure the projection matches the TIF files
transformer = Transformer.from_crs("epsg:4326", "epsg:4326")

# Function to extract climate data from a TIF file at a specific lat/lon
def extract_var (tif_file, lat, lon):
    try:
        with rasterio.open(tif_file) as dataset:
            # Transform coordinates if needed
            lat, lon = transformer.transform(lat, lon)
            # Get the row and column of the pixel corresponding to the lat/lon
            row, col = dataset.index(lon, lat)
            # Read the value of the pixel
            value = dataset.read(1)[row, col]
        return value
    
    except Exception as e:
        print(f"Error processing {tif_file} at lat: {lat}, lon: {lon}")
        print(e)
        return None

# List to store climate variable names
climate_vars = []

# Process each TIF file in the climate directory
for tif_file in os.listdir(climate_dir):
    if tif_file.endswith('.tif'):
        climate_var = tif_file.replace('.tif', '')
        climate_vars.append(climate_var)
        climate_file = os.path.join(climate_dir, tif_file)
        print(tif_file)
        
        # Extract climate data for each plot
        FIAdata[climate_var] = FIAdata.apply(lambda row: extract_var(climate_file, row['LAT'], 
                                                                     row['LON']), axis=1)

# Save the combined data to a new CSV file
output_file = os.path.join(parent_directory, 'NE_Plot_Composition_Climate.csv')
FIAdata.to_csv(output_file, index=False)

print(f"Data saved to {output_file}")